
#include "RSAEncrypt.h"

/** Construct the class with no bucket brigades. This means that this
 * is to be used for block and not pipe encryption.
 * 
 * @param rsapubkey The RSA public key to use for encryption. */
RSAEncrypt::RSAEncrypt(const RSAPublicKey &rsapubkey)
{
   rsapubkey.getSSLKeyComponents(&key);
   input = NULL;
   output = NULL;
   previous_block = NULL;
}

/** Construct the class with bucket brigades. This means that this can
 * be used for both block and pipe encryption.
 * 
 * @param iinput The bucket brigade for input.
 * 
 * @param ioutput The bucket brigade for output.
 * 
 * @param rsapubkey The RSA public key to use for encryption. */
RSAEncrypt::RSAEncrypt(APRBucketBrigade &iinput, APRBucketBrigade &ioutput, const RSAPublicKey &rsapubkey)
{
   rsapubkey.getSSLKeyComponents(&key);
   input = &iinput;
   output = &ioutput;
   previous_block = NULL;
}

/** Destructor for the RSA public key encryption class. */
RSAEncrypt::~RSAEncrypt()
{
   RSA_free(key);
}

/** Encrypt data block and return encrypted data. The encrypted data
 * returned must be freed by the caller.
 * 
 * @param plainblock A block of plain data to encrypt.
 * 
 * @param plainsize The length of the plain data block.
 * 
 * @param[out] codedsize The length of the coded block of data
 * returned. The data returned must be freed by the caller. */
char *RSAEncrypt::processDataBlock(const char *plainblock, apr_size_t plainsize, apr_size_t *codedsize)
{
   if (plainblock == NULL)
   {
      errprint("tried to encrypt NULL data block");
      return NULL;
   }
   if (key == NULL)
   {
      errprint("key is NULL - encryption failed");
      return NULL;
   }
   // For the PKCS v1.5 padding to be included too, apparently.
   apr_size_t plaindata_blocksze = RSA_size(key) - 12;
   apr_size_t codeddata_blocksze = RSA_size(key);
   int noofplaindatablocks = ((plaindata_blocksze + plainsize - 1) / (plaindata_blocksze));
   apr_size_t t_codedsize = noofplaindatablocks * codeddata_blocksze;
   char *returnblock = (char *)malloc(sizeof(char)*t_codedsize);
   previous_block = (char *)calloc(codeddata_blocksze, sizeof(char));
   char *manipdata = (char *)malloc(sizeof(char)*plaindata_blocksze);
   int j = 0;
   for (int i = 0; i < plainsize; i += plaindata_blocksze)
   {
      int tplainblock = ((plainsize - i) < plaindata_blocksze) ? (plainsize - i) : plaindata_blocksze;
      memcpy(manipdata, &(plainblock[i]), sizeof(char)*tplainblock);
      for (int k = 0; k < tplainblock; k++)
         manipdata[k] ^= previous_block[k];
      int thisblock = RSA_public_encrypt(tplainblock,
         (const unsigned char *)(manipdata),
         (unsigned char *)(&(returnblock[j])), key, RSA_PKCS1_PADDING);
      memcpy(previous_block, &(returnblock[j]), thisblock);
      j += thisblock;
   }
   if (j > t_codedsize)
   {
      errprint("overflow during encryption");
      return NULL;
   }
   free(manipdata);
   if (codedsize) *codedsize = j;
   free(previous_block);
   previous_block = NULL;
   return returnblock;
}

/** Begin the process of encrypting a pipe of data.
 * 
 * @return Zero on success, non-zero on failure. */
int RSAEncrypt::processPipeBegin()
{
   if (input == NULL)
   {
      errprint("no input pipe (is NULL)");
      return 1;
   }
   if (output == NULL)
   {
      errprint("no output pipe (is NULL)");
      return 1;
   }
   if (key == NULL)
   {
      errprint("key is NULL - encryption failed");
      return 1;
   }
   // In the RSA encryption there is no state, but we should allocate
   // the previous block
   previous_block = (char *)calloc(RSA_size(key), sizeof(char));
   return 0;
}

/** Process any outstanding plaintext data on the input pipe, and
 * place any processed encrypted data into the output pipe.
 * 
 * @return Zero on success, non-zero on failure. */
int RSAEncrypt::processPipeStep()
{
   // Padding
   apr_size_t plaindata_blocksize = RSA_size(key) - RSA_PADDING_SIZE;
   apr_size_t codeddata_blocksize = RSA_size(key);
   char *itemp_block = (char *)malloc(sizeof(char)*plaindata_blocksize);
   char *otemp_block = (char *)malloc(sizeof(char)*codeddata_blocksize);
   
   while (1)
   {
      apr_size_t readbytes = input->readFromStream(itemp_block, plaindata_blocksize);
      // this isn't a problem, this should fail during normal
      // operation when the buffer runs out of data.
      if (readbytes != plaindata_blocksize) break;
      for (int k = 0; k < plaindata_blocksize; k++)
         itemp_block[k] ^= previous_block[k];
      RSA_public_encrypt(plaindata_blocksize,
         (const unsigned char *)itemp_block, (unsigned char *)otemp_block, key, RSA_PKCS1_PADDING);
      memcpy(previous_block, otemp_block, codeddata_blocksize);
      apr_size_t writebytes = output->writeToStream(otemp_block, codeddata_blocksize);
      if (writebytes != codeddata_blocksize)
      {
         errprint("failed to write data to coded bucket brigade (wrote %" APR_SIZE_T_FMT ", had %" APR_SIZE_T_FMT "), likely corrupt", writebytes, codeddata_blocksize);
         free(itemp_block);
         free(otemp_block);
         return 1;
      }
      input->removeFromStream(plaindata_blocksize);
   }
   free(itemp_block);
   free(otemp_block);
   return 0;
}

/** Process the pipe step and record the movement of data to determine
 * when the process is complete.
 * 
 * @param[out] movementin Records data input in bytes.
 * 
 * @param[out] movementout Records data output in bytes. */
int RSAEncrypt::processPipeStep(int *movementin, int *movementout)
{
   apr_size_t before_input = input->realLength();
   apr_size_t before_output = output->realLength();
   int retval = processPipeStep();
   apr_size_t after_input = input->realLength();
   apr_size_t after_output = output->realLength();
   if (movementin)
      *movementin = before_input - after_input;
   if (movementout)
      *movementout = after_output - before_output;
   return retval;
}

/** Process any outstanding plaintext data in the pipe and finalise
 * the encryption process.
 * 
 * @return Zero on success, non-zero on failure. */
int RSAEncrypt::processPipeEnd()
{
   processPipeStep();
   // Padding
   apr_size_t plaindata_blocksize = input->realLength();
   if (plaindata_blocksize >= (RSA_size(key) - RSA_PADDING_SIZE))
   {
      errprint("too much data remaining in bucket brigade (is it still being written to?)");
      return 1;
   }
   apr_size_t codeddata_blocksize = RSA_size(key);
   char *itemp_block = (char *)malloc(sizeof(char)*plaindata_blocksize);
   char *otemp_block = (char *)malloc(sizeof(char)*codeddata_blocksize);
   apr_size_t readbytes = input->readFromStream(itemp_block, plaindata_blocksize);
   // this isn't a problem, this should fail during normal operation
   // when the buffer runs out of data.
   if (readbytes != plaindata_blocksize)
   {
      errprint("failed to read data from plain bucket brigade (read %" APR_SIZE_T_FMT ", had %" APR_SIZE_T_FMT "), buffer unfinished", readbytes, plaindata_blocksize);
      free(itemp_block);
      free(otemp_block);
      return 1;
   }
   for (int k = 0; k < plaindata_blocksize; k++)
      itemp_block[k] ^= previous_block[k];
   RSA_public_encrypt(plaindata_blocksize,
      (const unsigned char *)itemp_block, (unsigned char *)otemp_block, key, RSA_PKCS1_PADDING);
   memcpy(previous_block, otemp_block, codeddata_blocksize);
   apr_size_t writebytes = output->writeToStream(otemp_block, codeddata_blocksize);
   if (writebytes != codeddata_blocksize)
   {
      errprint("failed to write data to coded bucket brigade (wrote %" APR_SIZE_T_FMT ", had %" APR_SIZE_T_FMT "), likely corrupt", writebytes, codeddata_blocksize);
      free(itemp_block);
      free(otemp_block);
      return 1;
   }
   input->removeFromStream(plaindata_blocksize);
   free(itemp_block);
   free(otemp_block);
   free(previous_block);
   previous_block = NULL;
   return 0;
}
